#ifndef LOGGING_CONTROLLER_H
#define LOGGING_CONTROLLER_H

#include "Controller.h"

/** @brief A plugin class that allows saving/replaying low-level commands
 * from disk.
 *
 * Use the boolean flags 'save' and 'replay' to control the logging. 
 * 'Replay' overrides the base controller's normal update function.
 * Be careful not to interleave the replay mode with commands to the base
 * controller!
 *
 * If 'onlyJointCommands' is true, only the joint commands qdes, dqdes,
 * torque, and desiredVelocity are replayed.
 * The standard servo parameters are left untouched.
 */
template <class BaseController>
class LoggingController : public BaseController
{
 public:
  LoggingController(Robot& robot);
  bool SaveLog(const char* fn) const;
  bool LoadLog(const char* fn);
  virtual void Update(Real dt);

  //helpers that compress the logs
  bool EqualCommand(const ActuatorCommand& a,const ActuatorCommand& b) const;
  bool EqualCommand(const RobotCommand& a,const RobotCommand& b) const;
  void RemoveDelays(Real maxDelayTime);

  bool save,replay;
  bool onlyJointCommands; 
  vector<pair<Real,RobotCommand> > trajectory;
  int replayIndex;
};






template <class BaseController>
LoggingController<BaseController>::LoggingController(Robot& robot)
  : BaseController(robot),save(false),replay(false),onlyJointCommands(false),replayIndex(0)
{}

template <class BaseController>
bool LoggingController<BaseController>::SaveLog(const char* fn) const
{
  File f;
  if(!f.Open(fn,FILEWRITE)) return false;
  int size = trajectory.size();
  if(!WriteFile(f,size)) return false;
  for(size_t i=0;i<trajectory.size();i++) {
    if(!WriteFile(f,trajectory[i].first)) return false;
    if(!WriteFile(f,trajectory[i].second)) return false;
  }
  f.Close();
  return true;
}

template <class BaseController>
bool LoggingController<BaseController>::LoadLog(const char* fn)
{
  File f;
  if(!f.Open(fn,FILEREAD)) return false;
  int size;
  if(!ReadFile(f,size)) return false;
  if(size < 0) return false;
  trajectory.resize(size);
  for(size_t i=0;i<trajectory.size();i++) {
    if(!ReadFile(f,trajectory[i].first)) return false;
    if(!ReadFile(f,trajectory[i].second)) return false;
  }
  f.Close();
  return true;
}

template <class BaseController>
void LoggingController<BaseController>::Update(Real dt)
{
  if(replay) {   //replay mode
    BaseController::time += dt;
    if(!trajectory.empty()) {
      //look up the right trajectory
      Assert(replayIndex < (int)trajectory.size());
      //go backwards
      while(trajectory[replayIndex].first > BaseController::time && replayIndex > 0) {
	replayIndex--;
      }
      //go forwards
      while(replayIndex+1 < (int) trajectory.size() &&
	    BaseController::time >= trajectory[replayIndex+1].first) {
	replayIndex++;
      }
      //printf("Replay time %g, index %d\n",BaseController::time,replayIndex);
      //read it out
      const RobotCommand& logCmd = trajectory[replayIndex].second;
      RobotCommand* actualCmd = BaseController::command;
      if(onlyJointCommands) {
	for(size_t i=0;i<actualCmd->actuators.size();i++) {
	  actualCmd->actuators[i].qdes = logCmd.actuators[i].qdes;
	  actualCmd->actuators[i].dqdes = logCmd.actuators[i].dqdes;
	  actualCmd->actuators[i].torque = logCmd.actuators[i].torque;
	  actualCmd->actuators[i].desiredVelocity = logCmd.actuators[i].desiredVelocity;
	}
      }
      else {
	*actualCmd = logCmd;
      }
    }
  }
  else {  //normal mode
    BaseController::Update(dt);

    if(save) {
      if(trajectory.empty() || !EqualCommand(trajectory.back().second,*BaseController::command))
	trajectory.push_back(pair<Real,RobotCommand>(BaseController::time,*BaseController::command));
    }
  }
}

template <class BaseController>
bool LoggingController<BaseController>::EqualCommand(const ActuatorCommand& a,const ActuatorCommand& b) const
{
  if(onlyJointCommands) {
    if(a.qdes != b.qdes) return false;
    if(a.dqdes != b.dqdes) return false;
    //TODO: does feedforward torque matter?
    if(a.torque != b.torque) return false;
    if(a.desiredVelocity != b.desiredVelocity) return false;
  }
  else {
    if(a.mode!=b.mode) return false;
    if(a.measureAngleAbsolute!=b.measureAngleAbsolute) return false;
    if(a.kP!=b.kP) return false;
    if(a.kI!=b.kI) return false;
    if(a.kD!=b.kD) return false;
    if(a.iterm!=b.iterm) return false;
    if(a.qdes!=b.qdes) return false;
    if(a.dqdes!=b.dqdes) return false;
    if(a.mode!=b.mode) return false;
    if(a.torque!=b.torque) return false;
    if(a.desiredVelocity!=b.desiredVelocity) return false;
  }
  return true;
}

template <class BaseController>
bool LoggingController<BaseController>::EqualCommand(const RobotCommand& a,const RobotCommand& b) const
{
  if(a.actuators.size() != b.actuators.size()) return false;
  for(size_t i=0;i<a.actuators.size();i++)
    if(!EqualCommand(a.actuators[i],b.actuators[i])) return false;
  return true;
}

template <class BaseController>
void LoggingController<BaseController>::RemoveDelays(Real maxDelayTime)
{
  Assert(replayIndex == 0);
  int lastEraseIndex = -1;
  for(size_t i=1;i<trajectory.size();i++) {
    if(EqualCommand(trajectory[i].second,trajectory[i-1].second)) {
      if(lastEraseIndex < 0) lastEraseIndex = i;
    }
    else {
      if(lastEraseIndex >= 0) {
	//printf("Erasing trajectory commands %d-%d\n",lastEraseIndex,i-1);
	trajectory.erase(trajectory.begin()+lastEraseIndex,trajectory.begin()+i);
	i = lastEraseIndex-1;
	lastEraseIndex = -1;
      }
    }
  }
  Real shift=0;
  for(size_t i=0;i<trajectory.size();i++) {
    trajectory[i].first -= shift;
    if(i > 0) {
      if(trajectory[i].first-trajectory[i-1].first > maxDelayTime) {
	//printf("Found delay of %g at %d\n",trajectory[i].first-trajectory[i-1].first,i);
	shift += trajectory[i].first-trajectory[i-1].first-maxDelayTime;
	//printf("New shift %g\n",shift);
	trajectory[i].first = trajectory[i-1].first+maxDelayTime;
      }
    }
  }
  //printf("Trimmed %g seconds from trajectory\n",shift);
}

#endif
